import logging
from pinecone import Pinecone, ServerlessSpec
from sklearn.datasets import load_digits
import numpy as np
from collections import Counter
from tqdm import tqdm

# 设置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

# 初始化Pinecone客户端
def initialize_pinecone(api_key):
    pinecone = Pinecone(api_key=api_key)
    return pinecone

# 管理索引：检查索引是否存在，如果不存在则创建新索引
def manage_index(pinecone, index_name):
    existing_indexes = pinecone.list_indexes()
    if any(index['name'] == index_name for index in existing_indexes):  # 索引存在
        logging.info(f"索引 '%s' 已存在", index_name)
    else:  # 索引不存在
        logging.info(f"索引 '%s' 不存在", index_name)
        logging.info(f"正在创建新索引 '%s'", index_name)
        pinecone.create_index(
            name=index_name,
            dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
            metric="euclidean",  # 使用欧氏距离
            spec=ServerlessSpec(cloud="aws", region="us-east-1")  # 服务器规格
        )
        logging.info(f"索引 '%s' 创建成功", index_name)
    return pinecone.Index(index_name)

# 加载数据集并准备向量数据
def load_and_prepare_data(train_test_split=0.8):
    digits = load_digits(n_class=10)
    X = digits.data
    y = digits.target
    total_samples = len(X)
    train_samples = int(total_samples * train_test_split)
    train_X = X[:train_samples]
    train_y = y[:train_samples]
    test_X = X[train_samples:]
    test_y = y[train_samples:]
    vectors = []
    for i, (image, label) in enumerate(zip(train_X, train_y)):
        vector_id = str(i)
        vector_values = image.tolist()
        metadata = {"label": int(label)}
        vectors.append((vector_id, vector_values, metadata))
    return vectors, test_X, test_y

def upload_data(index, vectors, batch_size=1000):
    for i in tqdm(range(0, len(vectors), batch_size), desc="上传数据进度"):
        batch = vectors[i:i + batch_size]
        index.upsert(batch)
    logging.info("成功上传了%d条数据到Pinecone", len(vectors))

# 在Pinecone索引中执行查询并预测结果
def query_and_predict(index, test_X, test_y, top_k=11):
    correct_predictions = 0
    for i in tqdm(range(len(test_X)), desc="测试k=11的准确率"):
        query_data = test_X[i].tolist()
        results = index.query(vector=query_data, top_k=top_k, include_metadata=True)
        labels = [match['metadata']['label'] for match in results['matches']]
        final_prediction = Counter(labels).most_common(1)[0][0] if labels else None
        if final_prediction == test_y[i]:
            correct_predictions += 1
    accuracy = correct_predictions / len(test_X)
    logging.info("k=%d时的准确率为：%.4f", top_k, accuracy)
    return accuracy

# 主函数，程序入口点
def main(api_key, index_name):
    pinecone = initialize_pinecone(api_key)
    index = manage_index(pinecone, index_name)
    vectors, test_X, test_y = load_and_prepare_data(train_test_split=0.8)
    upload_data(index, vectors)
    accuracy = query_and_predict(index, test_X, test_y, top_k=11)
# 检查脚本是否作为主程序运行
if __name__ == "__main__":
    main("0f4bb63e-9728-43d4-89b8-562f96acb31b", "mnist-index")